from collections.abc import Iterator, Mapping
from functools import cache, cached_property, partial
from multiprocessing.pool import ThreadPool
from typing import Any, Callable, Sequence, TypeVar

from zkl_pyutils_fsspec import FsLike, resolve_child_fs

from zkl_aiutils_datasets.basics import Dataset
from zkl_aiutils_datasets.grouping import SimpleIndexedNamedDatasets
from zkl_aiutils_datasets.loading import load_dataset_by_scripting_from_itself
from zkl_aiutils_datasets.scripting import DatasetMetadata, load_dataset_metadata
from zkl_aiutils_datasets.utils.lazy import LazyMapping

AnySample = TypeVar('AnySample', bound=Any)
AnyDataset = TypeVar('AnyDataset', bound=Dataset)


class DirectoryDatasets(SimpleIndexedNamedDatasets[AnyDataset]):
    @classmethod
    def discover_dataset(cls,
        fs: FsLike, name: str,
        child_factory: Callable[[FsLike], AnyDataset],
    ) -> tuple[str, AnyDataset] | None:
        try:
            return name, child_factory(resolve_child_fs(fs, name))
        except Exception:
            return None

    @classmethod
    def _iter_discover_children(cls,
        fs: FsLike,
        child_factory: Callable[[FsLike], AnyDataset],
    ) -> Iterator[tuple[str, AnyDataset]]:
        with ThreadPool() as pool:
            children_name = fs.listdir("", detail=False)
            iterator = pool.imap_unordered(lambda name: cls.discover_dataset(fs, name, child_factory), children_name)
            for result in iterator:
                if result is not None:
                    yield result

    @classmethod
    def _discover_children(cls,
        fs: FsLike,
        child_factory: Callable[[FsLike], AnyDataset],
    ) -> tuple[tuple[str, AnyDataset], ...]:
        iterator = cls._iter_discover_children(fs, child_factory)
        return tuple(sorted(iterator, key=lambda x: x[0]))

    def __init__(self,
        fs: FsLike,
        child_factory: Callable[[FsLike], AnyDataset] = load_dataset_by_scripting_from_itself,
    ):
        self._fs = fs

        if self.children_metadata is None:
            super().__init__(self._discover_children(fs, child_factory))
        else:
            super().__init__(LazyMapping(
                (name, cache(partial(child_factory, resolve_child_fs(fs, name))))
                for name in self.children_metadata.keys()))

    @property
    def fs(self) -> str:
        return self._fs

    @cached_property
    def metadata(self) -> DatasetMetadata | None:
        try:
            return load_dataset_metadata(self._fs)
        except IOError:
            return None

    @cached_property
    def children_metadata(self) -> Mapping[str, DatasetMetadata] | None:
        metadata = self.metadata
        if metadata is None:
            return None
        # noinspection PyTypedDict
        children_metadata = metadata.get('children', None)
        if not isinstance(children_metadata, Mapping):
            return None
        return children_metadata

    @property
    def indexed_children(self) -> Sequence[AnyDataset]:
        return self._indexed_children

    @property
    def named_children(self) -> Mapping[str, AnyDataset]:
        return self._named_children
